/*
 * Copyright (c) Martin Kinkelin
 *
 * See the "License.txt" file in the root directory for infos
 * about permitted and prohibited uses of this code.
 */

using System;
using System.Windows.Forms;
using System.Diagnostics;
using System.Runtime.InteropServices;
using System.Threading;

namespace Jarf
{
	#region ShutdownType enum.

	/// <summary>
	/// Enumeration of supported shutdown types.
	/// </summary>
	public enum ShutdownType
	{
		/// <summary>
		/// Shuts down all processes running in the logon session of the current process.
		/// Then it logs the user off.
		/// This flag can be used only if the current process is running in an interactive user's logon session.
		/// </summary>
		LogOff = 0x0,
		/// <summary>
		/// Shuts down the system and turns off the power. The system must support the power-off feature.
		/// The current process must have the SE_SHUTDOWN_NAME privilege.
		/// </summary>
		PowerOff = 0x8,
		/// <summary>
		/// Shuts down the system and then restarts the system.
		/// The current process must have the SE_SHUTDOWN_NAME privilege.
		/// </summary>
		Reboot = 0x2,
		/// <summary>
		/// Suspends the system.
		/// </summary>
		Suspend = -1,
		/// <summary>
		/// Hibernates the system.
		/// </summary>
		Hibernate = -2,
	}

	#endregion

	/// <summary>
	/// Manages the shutdown of Windows.
	/// </summary>
	public class WindowsShutdownManager
	{
		#region WinAPI types, constants and functions.

		/// <summary>
		/// Please note: the packing value of 1 is absolutely necessary!
		/// </summary>
		[StructLayout(LayoutKind.Sequential, Pack = 1)]
		private struct TOKEN_PRIVILEGES
		{
			public uint PrivilegeCount;
			public long Luid;
			public uint Attributes;
		}



		private const uint TOKEN_ADJUST_PRIVILEGES = 0x20;

		private const uint SE_PRIVILEGE_ENABLED = 0x2;

		private const uint EWX_FORCE = 0x4;
		private const uint EWX_FORCEIFHUNG = 0x10;

		private const uint SHTDN_REASON_FLAG_PLANNED = 0x80000000;



		[DllImport("advapi32.dll")]
		private static extern int OpenProcessToken(IntPtr ProcessHandle, uint DesiredAccess, ref IntPtr TokenHandle);

		[DllImport("advapi32.dll")]
		private static extern int LookupPrivilegeValue(IntPtr lpSystemName, string lpName, ref long lpLuid);

		[DllImport("advapi32.dll", SetLastError = true)]
		private static extern int AdjustTokenPrivileges(IntPtr TokenHandle, int DisableAllPrivileges, ref TOKEN_PRIVILEGES NewState, uint BufferLength, IntPtr PreviousState, IntPtr ReturnLength);

		[DllImport("kernel32.dll")]
		private static extern int CloseHandle(IntPtr hObject);

		[DllImport("user32.dll", SetLastError = true)]
		private static extern int ExitWindowsEx(uint uFlags, uint dwReason);

		#endregion



		/// <summary>
		/// Tries to acquire the shutdown privilege for the current process and returns true if successful.
		/// </summary>
		/// <returns></returns>
		public static bool AcquireShutdownPrivilege()
		{
			TOKEN_PRIVILEGES newPrivileges = new TOKEN_PRIVILEGES();
			newPrivileges.PrivilegeCount = 1;
			newPrivileges.Luid = 0;
			newPrivileges.Attributes = SE_PRIVILEGE_ENABLED;

			IntPtr hToken = IntPtr.Zero;

			if (OpenProcessToken(Process.GetCurrentProcess().Handle, TOKEN_ADJUST_PRIVILEGES, ref hToken) == 0)
				return false;

			if (LookupPrivilegeValue(IntPtr.Zero, "SeShutdownPrivilege", ref newPrivileges.Luid) == 0 ||
			    AdjustTokenPrivileges(hToken, 0, ref newPrivileges, 0, IntPtr.Zero, IntPtr.Zero) == 0)
			{
				CloseHandle(hToken);
				return false;
			}

			CloseHandle(hToken);

			return (Marshal.GetLastWin32Error() == 0);
		}

		/// <summary>
		/// Tries to shut the system down and returns true if successful.
		/// </summary>
		/// <param name="type">Desired shutdown type.</param>
		/// <param name="force">
		/// Indicates whether applications which do not terminate in the time-out interval
		/// are to be killed.
		/// </param>
		/// <param name="planned">
		/// Indicates whether the shutdown is to be marked as planned.
		/// Ignored when suspending or hibernating the system.
		/// </param>
		/// <returns></returns>
		public static bool Shutdown(ShutdownType type, bool force, bool planned)
		{
			if (type == ShutdownType.Suspend)
				return Application.SetSuspendState(PowerState.Suspend, force, false);
			else if (type == ShutdownType.Hibernate)
				return Application.SetSuspendState(PowerState.Hibernate, force, false);
			else
			{
				uint flags = (uint)type;
				if (force)
					flags |= EWX_FORCEIFHUNG;

				return (ExitWindowsEx(flags, (planned ? SHTDN_REASON_FLAG_PLANNED : 0)) != 0);
			}
		}
	}
}
